package collections.map.tree_map;

import java.util.Comparator;
import java.util.ConcurrentModificationException;
import java.util.Iterator;
import java.util.NoSuchElementException;

import collections.map.MyMap;

public class MyTreeMap<K, V> implements MyMap<K, V> {

	public static final boolean BLACK = true;
	public static final boolean RED = false;
	
	private Comparator<? super K> comparator;
	private int size;
	private int modificationCount;
	private SimpleEntry<K, V> root;
	
	public MyTreeMap() {
		this.comparator = null;
	}
	
	public MyTreeMap(Comparator<K> comparator) {
		this.comparator = comparator;
	}

	@Override
	public void clear() {
		size = 0;
		root = null;
	}

	@Override
	public boolean containsKey(K key) {
		return get(key) != null;
	}

	@Override
	public boolean containsValue(V value) {
        for (SimpleEntry<K,V> e = getFirstEntry(); e != null; e = successor(e))
            if (valEquals(value, e.value))
                return true;
        return false;
	}

	private boolean valEquals(V value, V value2) {
        return (value == null ? value2 == null : value.equals(value2));
	}

	private SimpleEntry<K, V> successor(SimpleEntry<K, V> entry) {
        if (entry == null)
            return null;
        else if (entry.right != null) {
            SimpleEntry<K,V> rightEntry = entry.right;
            while (rightEntry.left != null)
                rightEntry = rightEntry.left;
            return rightEntry;
        } else {
            SimpleEntry<K,V> parentEntry = entry.parent;
            SimpleEntry<K,V> tempEntry = entry;
            while (parentEntry != null && tempEntry == parentEntry.right) {
                tempEntry = parentEntry;
                parentEntry = parentEntry.parent;
            }
            return parentEntry;
        }
	}
	
    private SimpleEntry<K, V> predecessor(SimpleEntry<K, V> entry) {
        if (entry == null)
            return null;
        else if (entry.left != null) {
            SimpleEntry<K,V> p = entry.left;
            while (p.right != null)
                p = p.right;
            return p;
        } else {
        	SimpleEntry<K,V> p = entry.parent;
        	SimpleEntry<K,V> ch = entry;
            while (p != null && ch == p.left) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
	}

	private SimpleEntry<K, V> getFirstEntry() {
        SimpleEntry<K,V> entry = root;
        if (entry != null)
            while (entry.left != null)
            	entry = entry.left;
        return entry;
	}

	@Override
    public V get(K key) {
        SimpleEntry<K,V> result = getEntry(key);
        return (result == null ? null : result.value);
    }

	private SimpleEntry<K, V> getEntry(K key) {
        if (comparator != null)
            return getEntryUsingComparator(key);
        if (key == null)
            throw new NullPointerException();
        Comparable<? super K> k = (Comparable<? super K>) key;
        SimpleEntry<K,V> result = root;
        while (result != null) {
            int cmp = k.compareTo(result.key);
            if (cmp < 0)
                result = result.left;
            else if (cmp > 0)
                result = result.right;
            else
                return result;
        }
        return null;
	}

	private SimpleEntry<K, V> getEntryUsingComparator(K key) {
        K k = (K) key;
        Comparator<? super K> comp = comparator;
        if (comp != null) {
            SimpleEntry<K,V> result = root;
            while (result != null) {
                int cmp = comp.compare(k, result.key);
                if (cmp < 0)
                    result = result.left;
                else if (cmp > 0)
                    result = result.right;
                else
                    return result;
            }
        }
        return null;
	}

	@Override
	public boolean isEmpty() {
		return size == 0;
	}

	@Override
	public V put(K key, V value) {
        SimpleEntry<K,V> t = root;
        if (t == null) {
            compare(key, key);

            root = new SimpleEntry<>(key, value, null);
            size = 1;
            modificationCount++;
            return null;
        }
        int cmp;
        SimpleEntry<K,V> parent;
        Comparator<? super K> cpr = comparator;
        if (cpr != null) {
            do {
                parent = t;
                cmp = cpr.compare(key, t.key);
                if (cmp < 0)
                    t = t.left;
                else if (cmp > 0)
                    t = t.right;
                else
                    return t.setValue(value);
            } while (t != null);
        }
        else {
            if (key == null)
                throw new NullPointerException();
            Comparable<? super K> k = (Comparable<? super K>) key;
            do {
                parent = t;
                cmp = k.compareTo(t.key);
                if (cmp < 0)
                    t = t.left;
                else if (cmp > 0)
                    t = t.right;
                else
                    return t.setValue(value);
            } while (t != null);
        }
        SimpleEntry<K,V> e = new SimpleEntry<>(key, value, parent);
        if (cmp < 0)
            parent.left = e;
        else
            parent.right = e;
        fixAfterInsertion(e);
        size++;
        modificationCount++;
        return null;
	}
	
    private void fixAfterInsertion(SimpleEntry<K, V> entry) {
        entry.color = RED;
        while (entry != null && entry != root && entry.parent.color == RED) {
            if (parentOf(entry) == leftOf(parentOf(parentOf(entry)))) {
                SimpleEntry<K,V> rightEntry = rightOf(parentOf(parentOf(entry)));
                if (colorOf(rightEntry) == RED) {
                    setColor(parentOf(entry), BLACK);
                    setColor(rightEntry, BLACK);
                    setColor(parentOf(parentOf(entry)), RED);
                    entry = parentOf(parentOf(entry));
                } else {
                    if (entry == rightOf(parentOf(entry))) {
                        entry = parentOf(entry);
                        rotateLeft(entry);
                    }
                    setColor(parentOf(entry), BLACK);
                    setColor(parentOf(parentOf(entry)), RED);
                    rotateRight(parentOf(parentOf(entry)));
                }
            } else {
                SimpleEntry<K,V> leftEntry = leftOf(parentOf(parentOf(entry)));
                if (colorOf(leftEntry) == RED) {
                    setColor(parentOf(entry), BLACK);
                    setColor(leftEntry, BLACK);
                    setColor(parentOf(parentOf(entry)), RED);
                    entry = parentOf(parentOf(entry));
                } else {
                    if (entry == leftOf(parentOf(entry))) {
                        entry = parentOf(entry);
                        rotateRight(entry);
                    }
                    setColor(parentOf(entry), BLACK);
                    setColor(parentOf(parentOf(entry)), RED);
                    rotateLeft(parentOf(parentOf(entry)));
                }
            }
        }
        root.color = BLACK;
	}

	private void rotateRight(SimpleEntry<K, V> entry) {
        if (entry != null) {
            SimpleEntry<K,V> leftEntry = entry.left;
            entry.left = leftEntry.right;
            if (leftEntry.right != null) leftEntry.right.parent = entry;
            leftEntry.parent = entry.parent;
            if (entry.parent == null)
                root = leftEntry;
            else if (entry.parent.right == entry)
                entry.parent.right = leftEntry;
            else entry.parent.left = leftEntry;
            leftEntry.right = entry;
            entry.parent = leftEntry;
        }
	}

	private void rotateLeft(SimpleEntry<K, V> entry) {
        if (entry != null) {
            SimpleEntry<K,V> rightEntry = entry.right;
            entry.right = rightEntry.left;
            if (rightEntry.left != null)
                rightEntry.left.parent = entry;
            rightEntry.parent = entry.parent;
            if (entry.parent == null)
                root = rightEntry;
            else if (entry.parent.left == entry)
                entry.parent.left = rightEntry;
            else
                entry.parent.right = rightEntry;
            rightEntry.left = entry;
            entry.parent = rightEntry;
        }
	}

	private SimpleEntry<K, V> rightOf(SimpleEntry<K, V> entry) {
        return (entry == null) ? null: entry.right;
	}

	private SimpleEntry<K, V> leftOf(SimpleEntry<K, V> entry) {
        return (entry == null) ? null: entry.left;
	}

	private void setColor(SimpleEntry<K, V> entry, boolean color) {
        if (entry != null)
            entry.color = color;
	}

	private boolean colorOf(SimpleEntry<K, V> entry) {
		return (entry == null ? BLACK : entry.color);
	}

	private SimpleEntry<K, V> parentOf(SimpleEntry<K, V> entry) {
		return (entry == null ? null: entry.parent);
	}

	final int compare(K key1, K key2) {
        return comparator==null ? ((Comparable<? super K>)key1).compareTo(key2)
            : comparator.compare(key1, key2);
    }

	@Override
	public V remove(K key) {
        SimpleEntry<K,V> entry = getEntry(key);
        if (entry == null)
            return null;

        V oldValue = entry.value;
        deleteEntry(entry);
        return oldValue;
	}

	private void deleteEntry(SimpleEntry<K, V> entry) {
        modificationCount++;
        size--;

        if (entry.left != null && entry.right != null) {
            SimpleEntry<K,V> s = successor(entry);
            entry.key = s.key;
            entry.value = s.value;
            entry = s;
        }

        SimpleEntry<K,V> replacement = (entry.left != null ? entry.left : entry.right);

        if (replacement != null) {
            replacement.parent = entry.parent;
            if (entry.parent == null)
                root = replacement;
            else if (entry == entry.parent.left)
                entry.parent.left  = replacement;
            else
                entry.parent.right = replacement;

            entry.left = entry.right = entry.parent = null;

            if (entry.color == BLACK)
                fixAfterDeletion(replacement);
        } else if (entry.parent == null) {
            root = null;
        } else {
            if (entry.color == BLACK)
                fixAfterDeletion(entry);

            if (entry.parent != null) {
                if (entry == entry.parent.left)
                    entry.parent.left = null;
                else if (entry == entry.parent.right)
                    entry.parent.right = null;
                entry.parent = null;
            }
        }		
	}

	private void fixAfterDeletion(SimpleEntry<K, V> entry) {
        while (entry != root && colorOf(entry) == BLACK) {
            if (entry == leftOf(parentOf(entry))) {
                SimpleEntry<K,V> sib = rightOf(parentOf(entry));

                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(entry), RED);
                    rotateLeft(parentOf(entry));
                    sib = rightOf(parentOf(entry));
                }

                if (colorOf(leftOf(sib))  == BLACK &&
                    colorOf(rightOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    entry = parentOf(entry);
                } else {
                    if (colorOf(rightOf(sib)) == BLACK) {
                        setColor(leftOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateRight(sib);
                        sib = rightOf(parentOf(entry));
                    }
                    setColor(sib, colorOf(parentOf(entry)));
                    setColor(parentOf(entry), BLACK);
                    setColor(rightOf(sib), BLACK);
                    rotateLeft(parentOf(entry));
                    entry = root;
                }
            } else { // symmetric
            	SimpleEntry<K,V> sib = leftOf(parentOf(entry));

                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(entry), RED);
                    rotateRight(parentOf(entry));
                    sib = leftOf(parentOf(entry));
                }

                if (colorOf(rightOf(sib)) == BLACK &&
                    colorOf(leftOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    entry = parentOf(entry);
                } else {
                    if (colorOf(leftOf(sib)) == BLACK) {
                        setColor(rightOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateLeft(sib);
                        sib = leftOf(parentOf(entry));
                    }
                    setColor(sib, colorOf(parentOf(entry)));
                    setColor(parentOf(entry), BLACK);
                    setColor(leftOf(sib), BLACK);
                    rotateRight(parentOf(entry));
                    entry = root;
                }
            }
        }

        setColor(entry, BLACK);
	}

	@Override
	public int size() {
		return size;
	}

	@Override
	public Iterator<Entry<K, V>> entryIterator() {
		return new EntryIterator(getFirstEntry());
	}
	
	private class EntryIterator implements Iterator<Entry<K, V>> {
		
        SimpleEntry<K,V> next;
        SimpleEntry<K,V> lastReturned;
        int expectedModCount;

        EntryIterator(SimpleEntry<K,V> first) {
            expectedModCount = modificationCount;
            lastReturned = null;
            next = first;
        }

        public boolean hasNext() {
            return next != null;
        }

        public SimpleEntry<K,V> next() {
        	SimpleEntry<K,V> entry = next;
            if (entry == null)
                throw new NoSuchElementException();
            if (modificationCount != expectedModCount)
                throw new ConcurrentModificationException();
            next = successor(entry);
            lastReturned = entry;
            return entry;
        }
        
		public void remove() {
            if (lastReturned == null)
                throw new IllegalStateException();
            if (modificationCount != expectedModCount)
                throw new ConcurrentModificationException();
            if (lastReturned.left != null && lastReturned.right != null)
                next = lastReturned;
            deleteEntry(lastReturned);
            expectedModCount = modificationCount;
            lastReturned = null;
        }
	}
	
	public class SimpleEntry<K, V> implements Entry<K, V> {

		K key;
		V value;
		SimpleEntry<K, V> parent;
		SimpleEntry<K, V> left = null;
		SimpleEntry<K, V> right = null;
		boolean color = BLACK;
		
		public SimpleEntry(K key, V value, SimpleEntry<K, V> parent) {
			this.key = key;
			this.value = value;
			this.parent = parent;
		}

		@Override
		public K getKey() {
			return key;
		}

		@Override
		public V getValue() {
			return value;
		}

		@Override
		public V setValue(V value) {
			V result = this.value;
			this.value = value;
			return result;
		}

		@Override
		public int hashCode() {
			final int prime = 31;
			int result = 1;
			result = prime * result + ((key == null) ? 0 : key.hashCode());
			result = prime * result + ((value == null) ? 0 : value.hashCode());
			return result;
		}

		@Override
		public boolean equals(Object obj) {
			if (this == obj)
				return true;
			if (obj == null)
				return false;
			if (getClass() != obj.getClass())
				return false;
			SimpleEntry other = (SimpleEntry) obj;
			if (key == null) {
				if (other.key != null)
					return false;
			} else if (!key.equals(other.key))
				return false;
			if (value == null) {
				if (other.value != null)
					return false;
			} else if (!value.equals(other.value))
				return false;
			return true;
		}
	}		
}
